import torch
import os
import csv
import time
from PIL import Image
from torchvision import transforms
from models.net import SimpleCNN
import config
from datetime import datetime
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np


class ImageEvaluator:
    def __init__(self, show_all_predictions=False, save_sample_image=False):
        self.device = config.DEVICE
        self.show_all_predictions = show_all_predictions
        self.save_sample_image = save_sample_image
        self.transform = transforms.Compose([
            transforms.Resize((28, 28)),
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
        self.model = self._load_model()
        self.output_dir = os.path.join(config.MODEL_SAVE_PATH, "evaluation_results")
        os.makedirs(self.output_dir, exist_ok=True)
        self.sample_image_dir = os.path.join(self.output_dir, "sample_images")
        os.makedirs(self.sample_image_dir, exist_ok=True)

    def _load_model(self):
        """Load the trained model"""
        model = SimpleCNN().to(self.device)
        model_path = os.path.join(config.MODEL_SAVE_PATH, "mnist_cnn_epoch_10.pth")
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model.eval()
        return model

    def _create_label_mapping(self, image_dir):
        """Create filename to label mapping based on naming convention: num_label.jpg or label_num.jpg"""
        label_mapping = {}
        valid_extensions = ('.jpg', '.jpeg', '.png')

        for f in os.listdir(image_dir):
            if f.lower().endswith(valid_extensions):
                try:
                    # 支持两种命名格式：
                    # 1. 序号_标签.扩展名 (如 13_2.jpeg)
                    # 2. 标签_序号.扩展名 (如 2_13.jpeg)
                    parts = f.split('_')
                    if len(parts) == 2:
                        # 尝试第一种格式：序号_标签
                        try:
                            label = int(parts[1].split('.')[0])  # 提取标签部分
                            if 0 <= label <= 9:
                                label_mapping[f] = label
                                continue
                        except (IndexError, ValueError):
                            pass

                        # 尝试第二种格式：标签_序号
                        try:
                            label = int(parts[0])  # 提取标签部分
                            if 0 <= label <= 9:
                                label_mapping[f] = label
                                continue
                        except (IndexError, ValueError):
                            pass

                    print(f"Warning: Filename '{f}' doesn't match expected naming conventions")
                except Exception as e:
                    print(f"Error processing filename '{f}': {str(e)}")
                    continue
        return label_mapping

    def _save_sample_image(self, img_tensor, pred_label, true_label, confidence, filename):
        """Save evaluation result as PNG image"""
        img = img_tensor.squeeze().cpu().numpy()
        img = (img * 0.3081) + 0.1307
        img = np.clip(img, 0, 1)

        plt.figure(figsize=(10, 5))

        plt.subplot(1, 2, 1)
        plt.imshow(img, cmap='gray')
        plt.title(f"Input Image\nTrue Label: {true_label}", fontsize=12)
        plt.axis('off')

        plt.subplot(1, 2, 2)
        with torch.no_grad():
            probs = F.softmax(self.model(img_tensor.to(self.device)), dim=1).cpu().numpy()[0]

        bars = plt.bar(range(10), probs, color='skyblue')
        bars[pred_label].set_color('red')
        if pred_label != true_label:
            bars[true_label].set_color('green')

        plt.xticks(range(10))
        plt.ylim(0, 1)
        plt.xlabel('Digit Class', fontsize=10)
        plt.ylabel('Probability', fontsize=10)
        plt.title('Prediction Probabilities\nPredicted: {} ({:.1f}%)\n{}'.format(
            pred_label,
            confidence * 100,
            "✓ Correct" if pred_label == true_label else "✗ Wrong"
        ), fontsize=12)
        plt.grid(True, linestyle='--', alpha=0.6)

        for rect in bars:
            height = rect.get_height()
            plt.text(rect.get_x() + rect.get_width() / 2., height,
                     f'{height:.2f}',
                     ha='center', va='bottom', fontsize=8)

        plt.tight_layout()

        save_path = os.path.join(self.sample_image_dir, f"eval_result_{filename.split('.')[0]}.png")
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close()
        print(f"Saved evaluation result to: {save_path}")

    def evaluate_images(self, image_dir):
        """Evaluate all images in the specified directory"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        csv_file = os.path.join(self.output_dir, f"eval_results_{timestamp}.csv")
        stats_file = os.path.join(self.output_dir, f"eval_stats_{timestamp}.txt")

        label_mapping = self._create_label_mapping(image_dir)
        if not label_mapping:
            print("Error: No valid labeled images found!")
            return

        image_files = [f for f in label_mapping.keys()]
        total_images = len(image_files)

        if total_images == 0:
            print(f"No properly labeled images found in directory: {image_dir}")
            return

        results = []
        start_time = time.time()
        correct_predictions = 0
        sample_saved = False

        with open(csv_file, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(["Image", "Predicted Label", "Confidence", "True Label", "Correct"])

            for idx, filename in enumerate(image_files):
                img_path = os.path.join(image_dir, filename)
                true_label = label_mapping[filename]

                try:
                    img = Image.open(img_path).convert('L')
                    tensor = self.transform(img).unsqueeze(0).to(self.device)

                    with torch.no_grad():
                        output = self.model(tensor)
                        pred_label = output.argmax().item()
                        confidence = F.softmax(output, dim=1)[0][pred_label].item()

                    is_correct = (pred_label == true_label)
                    if is_correct:
                        correct_predictions += 1

                    results.append({
                        'image': filename,
                        'image_path': img_path,
                        'pred': pred_label,
                        'confidence': confidence,
                        'true_label': true_label,
                        'correct': is_correct,
                        'tensor': tensor
                    })

                    writer.writerow([
                        filename,
                        pred_label,
                        f"{confidence:.4f}",
                        true_label,
                        "Yes" if is_correct else "No"
                    ])

                    if self.save_sample_image and not sample_saved:
                        self._save_sample_image(
                            tensor, pred_label, true_label, confidence,
                            filename
                        )
                        sample_saved = True

                except Exception as e:
                    print(f"Error processing {filename}: {str(e)}")
                    continue

        infer_time = time.time() - start_time
        accuracy = (correct_predictions / total_images * 100) if total_images > 0 else 0
        speed = total_images / infer_time if infer_time > 0 else 0

        with open(stats_file, 'w') as f:
            f.write("========== Evaluation Summary ==========\n")
            f.write(f"Evaluation Time: {timestamp}\n")
            f.write(f"Image Directory: {image_dir}\n")
            f.write(f"Total Images Processed: {total_images}\n")
            f.write(f"Correct Predictions: {correct_predictions}\n")
            f.write(f"Accuracy: {accuracy:.2f}%\n")
            f.write(f"Inference Speed: {speed:.2f} images/sec\n\n")

            f.write("========== Detailed Prediction Results ==========\n")
            f.write("Image Path | Image Name | True Label | Predicted Label | Confidence | Correct\n")
            f.write("-" * 90 + "\n")
            for result in results:
                f.write(f"{result['image_path']} | {result['image']} | {result['true_label']} | "
                        f"{result['pred']} | {result['confidence']:.4f} | "
                        f"{'Yes' if result['correct'] else 'No'}\n")

            incorrect = [r for r in results if not r['correct']]
            if incorrect:
                f.write("\n========== Incorrect Predictions Summary ==========\n")
                f.write(f"Total Incorrect: {len(incorrect)}\n")
                f.write("Image Name | True Label | Predicted Label | Confidence\n")
                f.write("-" * 60 + "\n")
                for item in incorrect:
                    f.write(f"{item['image']} | {item['true_label']} | {item['pred']} | {item['confidence']:.2%}\n")

            f.write("\nNote: Supported filename formats:\n")
            f.write("- [number]_[label].extension (e.g., 13_2.jpeg)\n")
            f.write("- [label]_[number].extension (e.g., 2_13.jpeg)\n")

        print(f"\nEvaluation Complete")
        print(f"Total images processed: {total_images}")
        print(f"Correct predictions: {correct_predictions}")
        print(f"Accuracy: {accuracy:.2f}%")
        print(f"Processing speed: {speed:.2f} images/sec")
        print(f"Detailed results saved to: {csv_file}")
        print(f"Complete statistics saved to: {stats_file}")
        if self.save_sample_image and sample_saved:
            print(f"Sample evaluation result saved to: {self.sample_image_dir}")


if __name__ == "__main__":
    evaluator = ImageEvaluator(show_all_predictions=False, save_sample_image=True)
    eval_dir = r"F:\pycharm_program\MNIST_project\mnist_project\data\evaluate_data"
    evaluator.evaluate_images(eval_dir)